/*
 * Copyright (c) 2024 OpenHiTLS. All rights reserved.
 * Licensed under the OpenHiTLS license.
 */
package org.openhitls.tls;

import javax.net.ssl.*;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.net.SocketAddress;
import java.nio.ByteBuffer;
import java.util.List;
import java.util.function.BiFunction;

/**
 * SSLSocket implementation backed by openHiTLS.
 *
 * <p>This class wraps a plain socket with TLS encryption using an internal
 * SSLEngine for the actual cryptographic operations.</p>
 */
public class HiTlsSSLSocket extends SSLSocket {

    private final HiTlsSSLContextSpi context;
    private HiTlsSSLEngine engine;
    private Socket underlyingSocket;
    private boolean autoClose;
    private boolean connected = false;
    private boolean handshakeCompleted = false;

    // I/O streams
    private InputStream sslInputStream;
    private OutputStream sslOutputStream;

    // Buffers for SSLEngine operations
    private ByteBuffer netInBuffer;
    private ByteBuffer netOutBuffer;
    private ByteBuffer appInBuffer;
    private ByteBuffer appOutBuffer;

    private static final int BUFFER_SIZE = 16384;

    /**
     * Creates an unconnected SSL socket.
     */
    HiTlsSSLSocket(HiTlsSSLContextSpi context) {
        this.context = context;
        this.underlyingSocket = new Socket();
        this.autoClose = true;
        initBuffers();
    }

    /**
     * Creates an SSL socket wrapping an existing socket.
     */
    HiTlsSSLSocket(HiTlsSSLContextSpi context, Socket socket, String host, int port, boolean autoClose)
            throws IOException {
        this.context = context;
        this.underlyingSocket = socket;
        this.autoClose = autoClose;
        this.connected = socket.isConnected();
        initBuffers();
        createEngine(host, port);
    }

    /**
     * Creates an SSL socket with consumed input data.
     */
    HiTlsSSLSocket(HiTlsSSLContextSpi context, Socket socket, InputStream consumed, boolean autoClose)
            throws IOException {
        this.context = context;
        this.underlyingSocket = socket;
        this.autoClose = autoClose;
        this.connected = socket.isConnected();
        initBuffers();
        // Handle consumed data if present
        if (consumed != null && consumed.available() > 0) {
            byte[] data = consumed.readAllBytes();
            netInBuffer.put(data);
            netInBuffer.flip();
        }
    }

    private void initBuffers() {
        netInBuffer = ByteBuffer.allocate(BUFFER_SIZE);
        netOutBuffer = ByteBuffer.allocate(BUFFER_SIZE);
        appInBuffer = ByteBuffer.allocate(BUFFER_SIZE);
        appOutBuffer = ByteBuffer.allocate(BUFFER_SIZE);
    }

    private void createEngine(String host, int port) {
        engine = (HiTlsSSLEngine) context.engineCreateSSLEngine(host, port);
        engine.setUseClientMode(true);
    }

    void connect(String host, int port) throws IOException {
        underlyingSocket.connect(new InetSocketAddress(host, port));
        createEngine(host, port);
        connected = true;
    }

    void connect(InetAddress address, int port) throws IOException {
        underlyingSocket.connect(new InetSocketAddress(address, port));
        createEngine(address.getHostAddress(), port);
        connected = true;
    }

    void bind(InetAddress localAddress, int localPort) throws IOException {
        underlyingSocket.bind(new InetSocketAddress(localAddress, localPort));
    }

    @Override
    public void connect(SocketAddress endpoint) throws IOException {
        connect(endpoint, 0);
    }

    @Override
    public void connect(SocketAddress endpoint, int timeout) throws IOException {
        InetSocketAddress addr = (InetSocketAddress) endpoint;
        underlyingSocket.connect(endpoint, timeout);
        createEngine(addr.getHostName(), addr.getPort());
        connected = true;
    }

    @Override
    public void startHandshake() throws IOException {
        if (!connected) {
            throw new IOException("Socket is not connected");
        }
        if (engine == null) {
            throw new IOException("SSLEngine not initialized");
        }

        engine.beginHandshake();

        // Perform handshake
        SSLEngineResult.HandshakeStatus status = engine.getHandshakeStatus();
        while (status != SSLEngineResult.HandshakeStatus.FINISHED &&
               status != SSLEngineResult.HandshakeStatus.NOT_HANDSHAKING) {

            switch (status) {
                case NEED_WRAP:
                    netOutBuffer.clear();
                    SSLEngineResult wrapResult = engine.wrap(appOutBuffer, netOutBuffer);
                    netOutBuffer.flip();

                    if (netOutBuffer.hasRemaining()) {
                        byte[] data = new byte[netOutBuffer.remaining()];
                        netOutBuffer.get(data);
                        underlyingSocket.getOutputStream().write(data);
                        underlyingSocket.getOutputStream().flush();
                    }
                    status = wrapResult.getHandshakeStatus();
                    break;

                case NEED_UNWRAP:
                    netInBuffer.clear();
                    byte[] received = new byte[BUFFER_SIZE];
                    int bytesRead = underlyingSocket.getInputStream().read(received);
                    if (bytesRead < 0) {
                        throw new SSLException("Connection closed during handshake");
                    }
                    netInBuffer.put(received, 0, bytesRead);
                    netInBuffer.flip();

                    appInBuffer.clear();
                    SSLEngineResult unwrapResult = engine.unwrap(netInBuffer, appInBuffer);
                    status = unwrapResult.getHandshakeStatus();
                    break;

                case NEED_TASK:
                    Runnable task;
                    while ((task = engine.getDelegatedTask()) != null) {
                        task.run();
                    }
                    status = engine.getHandshakeStatus();
                    break;

                default:
                    throw new SSLException("Unexpected handshake status: " + status);
            }
        }

        handshakeCompleted = true;
    }

    @Override
    public SSLSession getSession() {
        return engine != null ? engine.getSession() : null;
    }

    @Override
    public SSLSession getHandshakeSession() {
        return engine != null ? engine.getHandshakeSession() : null;
    }

    @Override
    public void addHandshakeCompletedListener(HandshakeCompletedListener listener) {
        // TODO: Implement listener support
    }

    @Override
    public void removeHandshakeCompletedListener(HandshakeCompletedListener listener) {
        // TODO: Implement listener support
    }

    @Override
    public String[] getSupportedCipherSuites() {
        return context.getSupportedCipherSuites();
    }

    @Override
    public String[] getEnabledCipherSuites() {
        return engine != null ? engine.getEnabledCipherSuites() : context.getEnabledCipherSuites();
    }

    @Override
    public void setEnabledCipherSuites(String[] suites) {
        if (engine != null) {
            engine.setEnabledCipherSuites(suites);
        }
    }

    @Override
    public String[] getSupportedProtocols() {
        return context.getSupportedProtocols();
    }

    @Override
    public String[] getEnabledProtocols() {
        return engine != null ? engine.getEnabledProtocols() : context.getEnabledProtocols();
    }

    @Override
    public void setEnabledProtocols(String[] protocols) {
        if (engine != null) {
            engine.setEnabledProtocols(protocols);
        }
    }

    @Override
    public void setUseClientMode(boolean mode) {
        if (engine != null) {
            engine.setUseClientMode(mode);
        }
    }

    @Override
    public boolean getUseClientMode() {
        return engine != null && engine.getUseClientMode();
    }

    @Override
    public void setNeedClientAuth(boolean need) {
        if (engine != null) {
            engine.setNeedClientAuth(need);
        }
    }

    @Override
    public boolean getNeedClientAuth() {
        return engine != null && engine.getNeedClientAuth();
    }

    @Override
    public void setWantClientAuth(boolean want) {
        if (engine != null) {
            engine.setWantClientAuth(want);
        }
    }

    @Override
    public boolean getWantClientAuth() {
        return engine != null && engine.getWantClientAuth();
    }

    @Override
    public void setEnableSessionCreation(boolean flag) {
        if (engine != null) {
            engine.setEnableSessionCreation(flag);
        }
    }

    @Override
    public boolean getEnableSessionCreation() {
        return engine == null || engine.getEnableSessionCreation();
    }

    @Override
    public SSLParameters getSSLParameters() {
        return engine != null ? engine.getSSLParameters() : new SSLParameters();
    }

    @Override
    public void setSSLParameters(SSLParameters params) {
        if (engine != null) {
            engine.setSSLParameters(params);
        }
    }

    @Override
    public InputStream getInputStream() throws IOException {
        if (sslInputStream == null) {
            sslInputStream = new SSLInputStream();
        }
        return sslInputStream;
    }

    @Override
    public OutputStream getOutputStream() throws IOException {
        if (sslOutputStream == null) {
            sslOutputStream = new SSLOutputStream();
        }
        return sslOutputStream;
    }

    @Override
    public void close() throws IOException {
        if (engine != null) {
            engine.closeOutbound();
        }
        if (autoClose && underlyingSocket != null) {
            underlyingSocket.close();
        }
    }

    @Override
    public boolean isClosed() {
        return underlyingSocket != null && underlyingSocket.isClosed();
    }

    @Override
    public boolean isConnected() {
        return connected && underlyingSocket != null && underlyingSocket.isConnected();
    }

    @Override
    public InetAddress getInetAddress() {
        return underlyingSocket != null ? underlyingSocket.getInetAddress() : null;
    }

    @Override
    public int getPort() {
        return underlyingSocket != null ? underlyingSocket.getPort() : -1;
    }

    @Override
    public InetAddress getLocalAddress() {
        return underlyingSocket != null ? underlyingSocket.getLocalAddress() : null;
    }

    @Override
    public int getLocalPort() {
        return underlyingSocket != null ? underlyingSocket.getLocalPort() : -1;
    }

    // ==================== Inner Stream Classes ====================

    private class SSLInputStream extends InputStream {
        @Override
        public int read() throws IOException {
            byte[] b = new byte[1];
            int n = read(b, 0, 1);
            return (n == -1) ? -1 : (b[0] & 0xFF);
        }

        @Override
        public int read(byte[] b, int off, int len) throws IOException {
            if (!handshakeCompleted) {
                startHandshake();
            }

            // Read from network
            netInBuffer.clear();
            byte[] received = new byte[BUFFER_SIZE];
            int bytesRead = underlyingSocket.getInputStream().read(received);
            if (bytesRead < 0) {
                return -1;
            }
            netInBuffer.put(received, 0, bytesRead);
            netInBuffer.flip();

            // Unwrap (decrypt)
            appInBuffer.clear();
            SSLEngineResult result = engine.unwrap(netInBuffer, appInBuffer);
            appInBuffer.flip();

            int available = appInBuffer.remaining();
            if (available == 0) {
                return 0;
            }

            int toRead = Math.min(len, available);
            appInBuffer.get(b, off, toRead);
            return toRead;
        }
    }

    private class SSLOutputStream extends OutputStream {
        @Override
        public void write(int b) throws IOException {
            write(new byte[]{(byte) b}, 0, 1);
        }

        @Override
        public void write(byte[] b, int off, int len) throws IOException {
            if (!handshakeCompleted) {
                startHandshake();
            }

            // Wrap (encrypt)
            appOutBuffer.clear();
            appOutBuffer.put(b, off, len);
            appOutBuffer.flip();

            netOutBuffer.clear();
            SSLEngineResult result = engine.wrap(appOutBuffer, netOutBuffer);
            netOutBuffer.flip();

            // Write to network
            if (netOutBuffer.hasRemaining()) {
                byte[] data = new byte[netOutBuffer.remaining()];
                netOutBuffer.get(data);
                underlyingSocket.getOutputStream().write(data);
            }
        }

        @Override
        public void flush() throws IOException {
            underlyingSocket.getOutputStream().flush();
        }
    }
}
